#include "include/rdMessage.h"

#include "util/Rijndael.h"
#include "util/whCRC.h"

#ifndef WIN32
#include <arpa/inet.h>
#else
#include <WinSock2.h>
#pragma comment(lib, "ws2_32.lib")
#endif

#include "core/thirdparty/zlib/zconf.h"
#include "core/thirdparty/zlib/zlib.h"

#define CHUNK 16384

namespace qb {
static bool s_crc = false;
static bool s_encrpt = false;
static bool s_zip = false;
static char AES_Key[32];
char tmp_Key[32] = "6&juWB0^^hG3uy7_";

rdMessage::rdMessage(char* buff, int len) {}

rdMessage::~rdMessage() {}

void rdMessage::Init(char* key) {
  int oLen = 0;
  memset(AES_Key, 0, 32);
  bDecryt(key, Y_RD_AES_KEY_LEN, AES_Key, &oLen, tmp_Key);
}

void rdMessage::SetPack(bool crc, bool encryt, bool zip) {
  s_crc = crc;
  s_encrpt = encryt;
  s_zip = zip;
}

bool rdMessage::bCheckCRC(char* in, int in_len, char* out, int* out_len) {
  unsigned int lCRC;
  unsigned int lMeCRC;
  memcpy(&lCRC, in + in_len - sizeof(int), sizeof(int));
  lCRC = ntohl(lCRC);
  CCRC oCRC(CCRC::E_CRC32);
  oCRC.CRCCheck(in, in_len - sizeof(int), &lMeCRC);
  if (lCRC == lMeCRC) {
    *out_len = in_len - sizeof(int);
    memcpy(out, in, *out_len);
    out[*out_len] = '\0';
    return true;
  } else {
    return false;
  }
}

bool rdMessage::bPackCRC(char* in, int in_len, unsigned int* out,
                         int* out_len) {
  CCRC oCRC(CCRC::E_CRC32);
  oCRC.CRCCheck(in, in_len, out);
  *out_len = sizeof(int);
  return true;
}

bool rdMessage::bEncryt(char* in, int in_len, char* out, int* out_len,
                        char* in_key) {
  CRijndael cCryptor;
  if (in_key == NULL)
    cCryptor.MakeKey(AES_Key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                     Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);
  else
    cCryptor.MakeKey(in_key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                     Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);

  if (in_len % Y_RD_AES_BLOCK_LEN == 0) {
    cCryptor.Encrypt(in, out, in_len, 1);
  } else {
    int ii = in_len / Y_RD_AES_BLOCK_LEN * Y_RD_AES_BLOCK_LEN;
    if (ii > 0) {
      cCryptor.Encrypt(in, out, ii, 1);
    }
    memcpy(out + ii, in + ii, in_len - ii);
  }

  *out_len = in_len;

  return true;
}

bool rdMessage::bDecryt(char* in, int in_len, char* out, int* out_len,
                        char* in_key) {
  CRijndael cCryptor;
  if (in_key == NULL)
    cCryptor.MakeKey(AES_Key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                     Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);
  else
    cCryptor.MakeKey(in_key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                     Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);

  if (in_len % Y_RD_AES_BLOCK_LEN == 0) {
    cCryptor.Decrypt(in, out, in_len, 1);
  } else {
    int ii = in_len / Y_RD_AES_BLOCK_LEN * Y_RD_AES_BLOCK_LEN;
    if (ii > 0) {
      cCryptor.Decrypt(in, out, ii, 1);
    }
    memcpy(out + ii, in + ii, in_len - ii);
  }
  *out_len = in_len;

  return true;
}

int rdMessage::iUnPack(char* in, int in_len, char* out, int* out_len) {
  unsigned char msgflag;
  bool crc, encryt, zip, order;
  int off = sizeof(unsigned char);
  memcpy(&msgflag, in, off);

  crc = msgflag & 0x01;
  encryt = msgflag & 0x02;
  zip = msgflag & 0x04;
  order = msgflag & 0x08;

  if (!crc && !encryt && !zip) {
    memcpy(out, in + off, in_len - off);
    *out_len = in_len - off;
    return 0;
  }
  int iDoLen = in_len - off;
  if (crc) {
    if (bCheckCRC(in + off, iDoLen, out, out_len) == false) {
      return -2;
    }
    iDoLen = *out_len;
  }
  if (order == false) {
    if (zip) {
      if (bUnZip(in + off, iDoLen, out, out_len) == false) {
        return -3;
      }
      memcpy(in + off, out, *out_len);
      iDoLen = *out_len;
    }
    if (encryt) {
      if (bDecryt(in + off, iDoLen, out, out_len) == false) {
        return -1;
      }
    }
  } else {
    if (encryt) {
      if (bDecryt(in + off, iDoLen, out, out_len) == false) {
        return -1;
      }
      memcpy(in + off, out, *out_len);
      iDoLen = *out_len;
    }
    if (zip) {
      if (bUnZip(in + off, iDoLen, out, out_len) == false) {
        return -3;
      }
      memcpy(in + off, out, *out_len);
      iDoLen = *out_len;
    }
  }
  return 0;
}

bool rdMessage::bPack(char* head, int h_len, char* body, int b_len, char* out,
                      int* out_len) {
  return bPack3(head, h_len, body, b_len, out, out_len);

  int magic = htonl(Y_RD_MSG_MAGIC);
  unsigned char msgflag = 0;
  int len = h_len + b_len + sizeof(msgflag);
  int off = 0;
  int len_pos = 0;
  int msg_pos = 0;

  if (s_crc) {
    msgflag = msgflag | 0x01;
    len += sizeof(int);
  }
  if (s_encrpt) {
    msgflag = msgflag | 0x02;
  }
  if (s_zip) {
    msgflag = msgflag | 0x04;
  }

  memcpy(out, &magic, sizeof(magic));
  off += sizeof(magic);
  len_pos = off;
  if (!s_crc && !s_encrpt && !s_zip) {
    len = htonl(len);
    memcpy(out + off, &len, sizeof(len));
    off += sizeof(len);
    memcpy(out + off, &msgflag, sizeof(msgflag));
    off += sizeof(msgflag);
    memcpy(out + off, head, h_len);
    off += h_len;
    memcpy(out + off, body, b_len);
    off += b_len;
    *out_len = off;
    return true;
  }
  // skip len
  off += sizeof(len);
  memcpy(out + off, &msgflag, sizeof(msgflag));
  off += sizeof(msgflag);
  msg_pos = off;

  char tcTemp[Y_RD_SERVICE_MESSAGE_BLOCK_LEN] = "";
  char Buf[Y_RD_SERVICE_MESSAGE_BLOCK_LEN] = "";
  int encryptlen = h_len + b_len;
  memcpy(Buf, head, h_len);
  memcpy(Buf + h_len, body, b_len);
  if (s_encrpt) {  // 加密
    if (bEncryt(Buf, h_len + b_len, tcTemp, &encryptlen) == false) {
      return false;
    }
    if (s_zip) {
      int ziplen = 0;
      if (!bZip(tcTemp, encryptlen, out + off, &ziplen)) return false;
      off += ziplen;
      len -= encryptlen - ziplen;
      len = htonl(len);
      memcpy(out + len_pos, &len, sizeof(len));  /// len_pos = 4?
      encryptlen = ziplen;
    } else {
      memcpy(out + len_pos, &len, sizeof(len));
      memcpy(out + off, tcTemp, encryptlen);
      off += encryptlen;
    }
  } else {  // 不加密
    if (s_zip) {
      int ziplen = 0;
      if (!bZip(Buf, len, out + off, &ziplen)) return false;
      off += ziplen;
      len -= encryptlen - ziplen;
      len = htonl(len);
      memcpy(out + len_pos, &len, sizeof(len));  /// len_pos = 4?
      encryptlen = ziplen;
    } else {
      memcpy(out + len_pos, &len, sizeof(len));
      memcpy(out + off, Buf, encryptlen);
      off += encryptlen;
    }
  }
  if (s_crc) {
    unsigned int lCRC;
    int iCRCLen;
    if (bPackCRC(out + msg_pos, encryptlen, &lCRC, &iCRCLen) == false) {
      return false;
    }
    lCRC = htonl(lCRC);
    memcpy(out + off, &lCRC, sizeof(int));
    off += iCRCLen;
  }
  *out_len = off;
  return true;
}

bool rdMessage::bPack2(char* head, int h_len, char* body, int b_len, char* out,
                       int* out_len) {
  char msg[Y_RD_SERVICE_MESSAGE_BLOCK_LEN] = {0};
  int off = 0;
  memset(msg, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
  memcpy(msg, head, h_len);
  memcpy(msg + h_len, body, b_len);

  int encryptlen = h_len + b_len;

  int len = h_len + b_len + 1;
  int len_pos = 0;
  int msg_pos = 0;
  unsigned char msgflag = 0;
  char sndBuf[Y_RD_SERVICE_MESSAGE_BLOCK_LEN] = {0};
  char tcTemp[Y_RD_SERVICE_MESSAGE_BLOCK_LEN] = {0};
  off = 0;
  memset(sndBuf, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);

  int magic = Y_RD_MSG_MAGIC;  // Y_WH_MSG_MAGIC;
  magic = htonl(magic);

  memcpy(sndBuf + off, &magic, sizeof(magic));
  off += sizeof(magic);  // 加magic
  len_pos = off;

  if (s_crc) {
    msgflag = msgflag | 0x01;
    len += sizeof(int);
  }
  if (s_encrpt) {
    msgflag = msgflag | 0x02;
  }
  if (s_zip) {
    msgflag = msgflag | 0x04;
  }

  int iLen = len;
  int iTotLen;

  iTotLen = iLen;
  iLen = htonl(iLen);
  memcpy(sndBuf + off, &iLen, sizeof(iLen));
  off += sizeof(iLen);  /// 加len
  memcpy(sndBuf + off, &msgflag, sizeof(msgflag));
  off += sizeof(msgflag);  // 加msgflag
  msg_pos = off;

  if (s_encrpt) {
    CRijndael m_Cryptor;
    m_Cryptor.MakeKey(AES_Key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                      Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);
    m_Cryptor.ResetChain();
    if (encryptlen % Y_RD_AES_BLOCK_LEN == 0) {
      memset(tcTemp, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
      m_Cryptor.Encrypt(msg, tcTemp, encryptlen, 1);
      if (s_zip) {
        int ziplen = 0;
        if (!bZip(tcTemp, encryptlen, sndBuf + off, &ziplen)) return 0;
        off += ziplen;
        iTotLen -= encryptlen - ziplen;
        iLen = htonl(iTotLen);  // 加msgflag长度
        memcpy(sndBuf + 4, &iLen, sizeof(iLen));
        encryptlen = ziplen;
      } else {
        memcpy(sndBuf + off, tcTemp, encryptlen);
        off += encryptlen;
      }
    } else {
      int ii = encryptlen / Y_RD_AES_BLOCK_LEN * Y_RD_AES_BLOCK_LEN;
      if (ii > 0) {
        memset(tcTemp, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
        m_Cryptor.Encrypt(msg, tcTemp, ii, 1);
      }
      memcpy(tcTemp + ii, msg + ii, encryptlen - ii);
      if (s_zip) {
        int ziplen = 0;
        if (!bZip(tcTemp, encryptlen, sndBuf + off, &ziplen)) return 0;
        off += ziplen;
        iTotLen -= encryptlen - ziplen;
        iLen = htonl(iTotLen);  // 加msgflag长度
        memcpy(sndBuf + 4, &iLen, sizeof(iLen));
        encryptlen = ziplen;
      } else {
        memcpy(sndBuf + off, tcTemp, encryptlen);
        off += encryptlen;
      }
    }
  } else {
    if (s_zip) {
      int ziplen = 0;
      if (!bZip(msg, encryptlen, sndBuf + off, &ziplen)) return 0;
      off += ziplen;
      iTotLen -= encryptlen - ziplen;
      iLen = htonl(iTotLen);  // 加msgflag长度
      memcpy(sndBuf + 4, &iLen, sizeof(iLen));
      encryptlen = ziplen;
    } else {
      memcpy(sndBuf + off, msg, encryptlen);
      off += encryptlen;
    }
  }
  if (s_crc) {
    unsigned int lCRC;
    int iCRCLen;
    if (bPackCRC(sndBuf + msg_pos /*sizeof(int) + sizeof(int)*/, encryptlen,
                 &lCRC, &iCRCLen) == false) {
      return false;
    }
    lCRC = htonl(lCRC);
    memcpy(sndBuf + off, &lCRC, sizeof(int));
    off += iCRCLen;
  }
  memcpy(out, sndBuf, off);
  *out_len = off;
  //	int iRet = m_SockStream.send_n(sndBuf,off);
  return true;
}

bool rdMessage::bPack3(char* head, int h_len, char* body, int b_len, char* out,
                       int* out_len) {
  char msg[Y_RD_SERVICE_MESSAGE_BLOCK_LEN];
  int off = 0;
  memset(msg, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
  memcpy(msg, head, h_len);
  if (b_len > 0) memcpy(msg + h_len, body, b_len);

  int encryptlen = h_len + b_len;

  int len = h_len + b_len + 1;
  int len_pos = 0;
  int msg_pos = 0;
  unsigned char msgflag = 0;
  char sndBuf[Y_RD_SERVICE_MESSAGE_BLOCK_LEN];
  char tcTemp[Y_RD_SERVICE_MESSAGE_BLOCK_LEN];
  char tcTempEncypt[Y_RD_SERVICE_MESSAGE_BLOCK_LEN];

  off = 0;
  memset(sndBuf, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);

  int magic = Y_RD_MSG_MAGIC;  // Y_WH_MSG_MAGIC;
  magic = htonl(magic);

  memcpy(sndBuf + off, &magic, sizeof(magic));
  off += sizeof(magic);  // 加magic
  len_pos = off;

  if (s_crc) {
    msgflag = msgflag | 0x01;
    len += sizeof(int);
  }
  if (s_encrpt) {
    msgflag = msgflag | 0x02;
  }
  if (s_zip) {
    msgflag = msgflag | 0x04;
  }

  // 区分于bPack2，先压缩后加密
  msgflag = msgflag | 0x08;

  int iLen = len;
  int iTotLen;

  iTotLen = iLen;
  iLen = htonl(iLen);
  memcpy(sndBuf + off, &iLen, sizeof(iLen));
  off += sizeof(iLen);  /// 加len
  memcpy(sndBuf + off, &msgflag, sizeof(msgflag));
  off += sizeof(msgflag);  // 加msgflag
  msg_pos = off;

  //////////////////
  memset(tcTemp, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
  if (s_zip) {
    int ziplen = 0;
    if (!bZip(msg, encryptlen, tcTemp, &ziplen)) return 0;
    // off += ziplen;
    iTotLen -= encryptlen - ziplen;
    iLen = htonl(iTotLen);  // 加msgflag长度
    memcpy(sndBuf + 4, &iLen, sizeof(iLen));
    encryptlen = ziplen;
  } else {
    memcpy(tcTemp, msg, encryptlen);
  }

  ///////////
  if (s_encrpt) {
    CRijndael m_Cryptor;
    m_Cryptor.MakeKey(AES_Key, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0",
                      Y_RD_AES_KEY_LEN, Y_RD_AES_BLOCK_LEN);
    m_Cryptor.ResetChain();
    if (encryptlen % Y_RD_AES_BLOCK_LEN == 0) {
      memset(tcTempEncypt, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
      m_Cryptor.Encrypt(tcTemp, tcTempEncypt, encryptlen, 1);
      memcpy(sndBuf + off, tcTempEncypt, encryptlen);
      off += encryptlen;
    } else {
      int ii = encryptlen / Y_RD_AES_BLOCK_LEN * Y_RD_AES_BLOCK_LEN;
      if (ii > 0) {
        memset(tcTempEncypt, 0x00, Y_RD_SERVICE_MESSAGE_BLOCK_LEN);
        m_Cryptor.Encrypt(tcTemp, tcTempEncypt, ii, 1);
      }
      memcpy(tcTempEncypt + ii, tcTemp + ii, encryptlen - ii);
      memcpy(sndBuf + off, tcTempEncypt, encryptlen);
      off += encryptlen;
    }
    if (off > Y_RD_SERVICE_MESSAGE_BLOCK_LEN) return false;
  } else {
    memcpy(sndBuf + off, tcTemp, encryptlen);
    off += encryptlen;
  }
  if (s_crc) {
    unsigned int lCRC;
    int iCRCLen;
    if (bPackCRC(sndBuf + msg_pos /*sizeof(int) + sizeof(int)*/, encryptlen,
                 &lCRC, &iCRCLen) == false) {
      return false;
    }
    lCRC = htonl(lCRC);
    if (off + iCRCLen > Y_RD_SERVICE_MESSAGE_BLOCK_LEN) return false;
    memcpy(sndBuf + off, &lCRC, sizeof(int));
    off += iCRCLen;
  }
  if (off > Y_RD_SERVICE_MESSAGE_BLOCK_LEN) return false;
  memcpy(out, sndBuf, off);
  *out_len = off;
  //	int iRet = m_SockStream.send_n(sndBuf,off);
  return true;
}

bool rdMessage::bZip(char* in, int in_len, char* out, int* out_len) {
  //	compress((unsigned char*)out, (unsigned long*)out_len, (unsigned
  //char*)in, in_len);

  //	return true;
  z_stream strm;
  int ret;

  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  // try
  {
    ret = deflateInit(&strm, Z_DEFAULT_COMPRESSION);
    if (ret != Z_OK) {
      return false;
    }
    strm.avail_in = in_len;
    strm.next_in = (Bytef*)in;
    strm.avail_out = Y_RD_SERVICE_MESSAGE_BLOCK_LEN;
    strm.next_out = (Bytef*)out;
    ret = deflate(&strm, Z_FINISH);
    if (ret != Z_STREAM_ERROR) {
      *out_len = Y_RD_SERVICE_MESSAGE_BLOCK_LEN - strm.avail_out;
      deflateEnd(&strm);
      return true;
    } else {
      deflateEnd(&strm);
      return false;
    }
  }
  // catch (...)
  {
    //	deflateEnd(&strm);
    //	return false;
  }
}

bool rdMessage::bUnZip(char* in, int in_len, char* out, int* out_len) {
  //	uncompress((unsigned char*)out, (unsigned long*)out_len, (unsigned
  //char*)in, in_len);

  //	return true;
  z_stream strm;
  int ret;

  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  strm.avail_in = 0;
  strm.next_in = Z_NULL;
  // try
  {
    ret = inflateInit(&strm);
    if (ret != Z_OK) {
      return false;
    }
    strm.avail_in = in_len;
    strm.next_in = (Bytef*)in;
    strm.avail_out = Y_RD_SERVICE_MESSAGE_BLOCK_LEN;
    strm.next_out = (Bytef*)out;
    ret = inflate(&strm, Z_FINISH);
    if (ret == Z_DATA_ERROR || ret == Z_NEED_DICT || ret == Z_MEM_ERROR) {
      inflateEnd(&strm);
      return false;
    } else {
      *out_len = Y_RD_SERVICE_MESSAGE_BLOCK_LEN - strm.avail_out;
      inflateEnd(&strm);
      return true;
    }
  }
  // catch (...)
  {
    //	inflateEnd(&strm);
    //	return false;
  }
}

bool rdMessage::UnCompress(char* in, int in_len, char* out, int* out_len,
                           int totallen) {
  z_stream strm;
  int ret;

  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;
  strm.avail_in = 0;
  strm.next_in = Z_NULL;
  // try
  {
    ret = inflateInit(&strm);
    if (ret != Z_OK) {
      return false;
    }
    strm.avail_in = in_len;
    strm.next_in = (Bytef*)in;
    strm.avail_out = totallen;
    strm.next_out = (Bytef*)out;
    ret = inflate(&strm, Z_FINISH);
    if (ret == Z_DATA_ERROR || ret == Z_NEED_DICT || ret == Z_MEM_ERROR) {
      inflateEnd(&strm);
      return false;
    } else {
      *out_len = totallen - strm.avail_out;
      inflateEnd(&strm);
      return true;
    }
  }
  // catch (...)
  {
    //	inflateEnd(&strm);
    //	return false;
  }

  return true;
}

}  // namespace qb
